from utils import print_timing, is_blankline
from string import whitespace

class Pheno_Base(object):
    def __init__(self, phenofiles, dbstring, table_dict):
        self.phenofiles = phenofiles
        self.dbstring = dbstring
        self.table_dict = table_dict

    def clinical_indices(self, phenofile):
        df = open(phenofile)
        header_split = df.readline().strip(whitespace).split(',')
        discard_chars = whitespace + '"'
        file_headers = [h.strip(discard_chars) for h in header_split]
        sfh = set(file_headers)
        headers = ["expid", "famid", "patientid", "phenotype", "race", "sex",  "studyid"]
        required_headers = ["expid", "sex"]
        sh = set(headers)
        srh = set(required_headers)
        # check that headers has no duplicated entries
        if len(set(file_headers)) != len(file_headers):
            print "abort: headers of %s contain duplicates"%phenofile
            exit(1)
        # check that header_list is contained in headers
        #
        if not srh.issubset(sfh):
            print "abort: headers of %s do not contain required headers %s"%(phenofile, list(srh.difference(sfh)))
            exit(1)
        index_dict = {}
        for h in sh.intersection(sfh):
            index_dict[h] = file_headers.index(h)
        df.close()
        return index_dict

    def get_patient_expids(self, phenofile):
        df = open(phenofile)
        df.readline()
        index = self.clinical_indices(phenofile)
        expid_index = index["expid"]
        patient_expids = []
        for line in df:
            if is_blankline(line):
                continue
            line_split = line.strip(whitespace).split(',')
            line_split = [x.strip(whitespace) for x in line_split]
            patient_expids.append(line_split[expid_index])
        return patient_expids

    def get_sex_dict(self, session):
        from dbschema import Sex
        sexd = {}
        sexd['1'] = Sex('M', 1, 1)
        sexd['2'] = Sex('F', 2, 0)
        sexd['-1'] = Sex('U', '-1', 'NA')
        return sexd

    def get_race_dict(self, session):
        from dbschema import Race
        # Race objects.
        raced = {}
        from utils import get_conf
        conf = get_conf()
        racevals = conf["race_values"]
        for r in racevals:
            raced[r] = Race(r)
        # Extra entries for compatability with data files.
        #raced['NA'] = raced['Unknown']
        return raced

    def get_race_fn(self, phenofile, index, raced):
        def return_race_obj(phenofile, line):
            line_split = line.strip(whitespace).split(',')
            line_split = [x.strip(whitespace) for x in line_split]
            raceval = line_split[race_index]
            if raceval not in raced:
                print "Error: the value of race in line"
                print line.rstrip("\n")
                print "of pheno file %s"%phenofile
                print "is %s, which is not listed as a permitted value"%(raceval)
                print "If you want to allow it, add %s to the race_values list in the config"%(raceval)
                exit(1)
            return raced[raceval]
        def return_none(phenofile, line): None
        if "race" in index:
            race_index = index["race"]
            return return_race_obj
        else:
            return return_none

    def make_pheno_table(self, phenofile, sexd, raced):
        print "making pheno table..."
        from dbschema import Pheno
        df = open(phenofile)
        df.readline()
        index = self.clinical_indices(phenofile)
        if "patientid" in index:
            patientid_index = index["patientid"]
            patientid_fn = lambda line_split: line_split[patientid_index]
        else:
            print "patientid missing, using expid"
            expid_index = index["expid"]
            patientid_fn = lambda line_split: line_split[expid_index]
        sex_index = index["sex"]
        if "famid" in index:
            famid_index = index["famid"]
            famid_fn = lambda line_split: line_split[famid_index]
        else:
            print "famid missing, using patientid"
            famid_fn = lambda line_split: patientid_fn(line_split)
        if "phenotype" in index:
            pheno_index = index["phenotype"]
            pheno_fn = lambda line_split: line_split[pheno_index]
        else:
            pheno_fn = lambda line_split: 0
        race_fn = self.get_race_fn(phenofile, index, raced)
        patientlst = []
        phenolst = []
        phenodict = {}
        for line in df:
            if is_blankline(line):
                continue
            line_split = line.strip(whitespace).split(',')
            line_split = [x.strip(whitespace) for x in line_split]
            famid = famid_fn(line_split)
            patientid = patientid_fn(line_split)
            phenotype = pheno_fn(line_split)
            race = race_fn(phenofile, line)
            sexval = line_split[sex_index]
            if sexval not in sexd:
                print "Error: the value of sex in line"
                print line.rstrip("\n")
                print "of pheno file %s"%phenofile
                print "is %s. Only sex values of 1, 2 and -1 are allowed."%(sexval)
                import sys
                sys.exit(1)
            else:
                sex = sexd[line_split[sex_index]]
            if patientid not in phenodict:
                phenodict[patientid] = {'famid':famid, 'sex':sex, 'race':race, 'phenotype':phenotype}
                phenolst.append(Pheno(patientid, famid, sex, race, phenotype))
            # NB: Temporarily disable this check
            else: #check for consistency for good measure
                pass
                #assert phenodict[patientid]['sex'] == sex, "sex not consistent for patientid %s"%patientid
                #assert phenodict[patientid]['race'] == race, "race not consistent for patientid %s"%patientid
        return phenolst

    def dummy_patients(self, expids, patient_expids, sexd, raced):
        # For throughness, should check that no patients have expids not in the expid list.
        # Add dummy patients corresponding to expids which don't have an associated patient
        from dbschema import Pheno, Sex
        nonpatient_expids = sorted(list(set(expids).difference(set(patient_expids))))
        phenolst = []
        for i, expid in enumerate(nonpatient_expids):
            patientid = 'DUMMY%s'%i
            sex = sexd['-1']
            race = raced['Unknown']
            phenolst.append(Pheno(patientid=patientid, famid=patientid, sex=sex, race=race))
        return phenolst

    def update_pheno_table(self):
        pdstr = ' and '.join(self.phenofiles)
        print "updating pheno table from %s..."%pdstr
        from utils import create_engine_wrapper
        pheno_table = self.table_dict['pheno_table']
        Pheno = self.table_dict['Pheno']
        meta = self.table_dict['metadata']
        from sqlalchemy.orm import sessionmaker
        db = create_engine_wrapper(self.dbstring)
        meta.bind = db
        meta.create_all(tables=[pheno_table])
        Session = sessionmaker()
        session = Session()
        sexd = self.get_sex_dict(session)
        raced = self.get_race_dict(session)
        phenolst = []
        for phenofile in self.phenofiles:
            phenolst = phenolst + self.make_pheno_table(phenofile, sexd, raced)
        # check ids in phenolst match those in the pheno table
        if sorted([p.patientid for p in phenolst]) != sorted([p.patientid for p in session.query(Pheno).all()]):
            print "pheno ids in the database and in %s differ, exiting"%pdstr
            exit(1)
        for p in phenolst:
            session.merge(p)
        session.commit()

    def make_idlink_table_pheno(self, phenofile, expids):
        #expids =  celfilenames(callsfilename)
        # Check that all expids are in genofile
        patient_expids = self.get_patient_expids(phenofile)
        diff = list(set(patient_expids).difference(set(expids)))
        if diff:
            print "WARNING: there are experimental ids in the pheno file %s which do not correspond to anything in the genotype calls file. These are %s"%(phenofile, diff)
            print "WARNING: the number of experimental ids in the pheno file %s which do not correspond to experimental ids in the genotype calls file is %s. See list above."%(phenofile, len(diff))
            print "WARNING: restricting loading to experimental ids in the genotype calls file."
            patient_expids = list(set(patient_expids).intersection(set(expids)))
        df = open(phenofile)
        df.readline()
        index = self.clinical_indices(phenofile)
        expid_index = index["expid"]
        if "patientid" in index:
            patientid_index = index["patientid"]
        else:
            print "patientid missing, using expid"
            patientid_index = index["expid"]
        if "studyid" in index:
            studyid_index = index["studyid"]
            studyid_fn = lambda line_split: line_split[studyid_index]
        else:
            print "studyid missing, using 'unknown'"
            studyid_fn = lambda line_split: 'unknown'
        insert_dict_list = []
        for line in df:
            if is_blankline(line):
                continue
            line_split = line.strip(whitespace).split(',')
            line_split = [x.strip(whitespace) for x in line_split]
            expid = line_split[expid_index]
            if expid in patient_expids:
                patientid = line_split[patientid_index]
                studyid = studyid_fn(line_split)
                insert_dict_list.append({'expid':expid,'patientid':patientid, 'studyid':studyid, 'sampleid':patientid})
        return insert_dict_list

    # create idlink table entries for genotype entries which don't corespond to patient entries.
    def make_idlink_table_dummy(self, expids, patient_expids):
        #expids =  celfilenames(callsfilename)
        nonpatient_expids = sorted(list(set(expids).difference(set(patient_expids))))
        insert_dict_list = []
        for i, expid in enumerate(nonpatient_expids):
            patientid = 'DUMMY%s'%i
            insert_dict_list.append({'expid':expid,'patientid':patientid, 'studyid':'unknown', 'sampleid':patientid})
        return insert_dict_list

    def make_idlink_table(self, db, expids, patient_expids):
        print "making idlink table..."
        idlink_table = self.table_dict['idlink_table']
        #expids =  celfilenames(callsfilename)
        conn = db.connect()
        insert_dict_list = []
        for phenofile in self.phenofiles:
            insert_dict_list.extend(self.make_idlink_table_pheno(phenofile, expids))
        insert_dict_list.extend(self.make_idlink_table_dummy(expids, patient_expids))
        insert_dict_list.sort(key=lambda(x):expids.index(x['expid']))
        conn.execute(idlink_table.insert(), insert_dict_list)
